"""
模型训练脚本
功能：训练语音情感识别模型（需要RAVDESS等数据集）
参数：数据集路径、模型类型等
返回值：训练好的模型文件
使用场景：离线训练模型，生成可部署的模型文件
运行方法：python scripts/train_model.py --data_path /path/to/dataset
"""

import argparse
import os
import sys
import pandas as pd
import numpy as np
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from features.audio_features import AudioFeatureExtractor
from models.emotion_classifier import EmotionClassifier
from core.config import settings

def parse_ravdess_filename(filename: str) -> str:
    """
    解析RAVDESS数据集文件名获取情感标签
    RAVDESS文件名格式：Modality-VocalChannel-Emotion-EmotionalIntensity-Statement-Repetition-Actor.wav
    参数：filename - RAVDESS文件名
    返回值：情感标签
    """
    # RAVDESS情感标签映射
    emotion_mapping = {
        '01': 'neutral',    # neutral
        '02': 'neutral',    # calm -> neutral  
        '03': 'happy',      # happy
        '04': 'sad',        # sad
        '05': 'angry',      # angry
        '06': 'fearful',    # fearful -> 暂不使用
        '07': 'disgust',    # disgust -> 暂不使用
        '08': 'surprised'   # surprised -> 暂不使用
    }
    
    try:
        # 分割文件名
        parts = filename.split('-')
        emotion_code = parts[2]
        
        # 返回映射的情感标签，过滤不支持的情感
        emotion = emotion_mapping.get(emotion_code)
        if emotion in settings.EMOTION_LABELS:
            return emotion
        else:
            return None  # 过滤不支持的情感
            
    except (IndexError, KeyError):
        return None

def load_dataset(data_path: str) -> tuple:
    """
    加载并预处理数据集
    参数：data_path - 数据集根目录路径
    返回值：(特征矩阵, 标签数组, 文件名列表)
    """
    feature_extractor = AudioFeatureExtractor()
    
    features_list = []
    labels_list = []
    files_list = []
    
    print("开始加载数据集...")
    
    # 遍历数据集目录
    for root, dirs, files in os.walk(data_path):
        for file in files:
            if file.endswith('.wav'):
                # 解析情感标签
                emotion = parse_ravdess_filename(file)
                
                if emotion is None:
                    continue  # 跳过不支持的情感或无法解析的文件
                
                file_path = os.path.join(root, file)
                
                try:
                    # 提取特征
                    features = feature_extractor.extract_all_features(file_path)
                    
                    features_list.append(features)
                    labels_list.append(emotion)
                    files_list.append(file)
                    
                    if len(features_list) % 50 == 0:
                        print(f"已处理 {len(features_list)} 个文件...")
                        
                except Exception as e:
                    print(f"处理文件 {file} 时出错: {e}")
                    continue
    
    if len(features_list) == 0:
        raise ValueError("未找到有效的音频文件，请检查数据集路径")
    
    # 转换为numpy数组
    features_array = np.array(features_list)
    labels_array = np.array(labels_list)
    
    print(f"数据集加载完成：")
    print(f"  - 总样本数: {len(features_list)}")
    print(f"  - 特征维度: {features_array.shape[1]}")
    print(f"  - 情感类别分布:")
    
    unique, counts = np.unique(labels_array, return_counts=True)
    for emotion, count in zip(unique, counts):
        print(f"    - {emotion}: {count} 样本")
    
    return features_array, labels_array, files_list

def plot_training_results(training_results: dict, save_path: str = "training_results.png"):
    """
    绘制训练结果图表
    参数：
        training_results - 训练结果字典
        save_path - 图表保存路径
    """
    plt.figure(figsize=(15, 5))
    
    # 混淆矩阵
    plt.subplot(1, 3, 1)
    cm = np.array(training_results['confusion_matrix'])
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=settings.EMOTION_LABELS,
                yticklabels=settings.EMOTION_LABELS)
    plt.title('混淆矩阵')
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')
    
    # 准确率对比
    plt.subplot(1, 3, 2)
    metrics = ['训练准确率', '验证准确率', '交叉验证均值']
    scores = [
        training_results['train_accuracy'],
        training_results['validation_accuracy'], 
        training_results['cv_mean_score']
    ]
    plt.bar(metrics, scores, color=['skyblue', 'lightgreen', 'orange'])
    plt.title('模型性能指标')
    plt.ylabel('准确率')
    plt.ylim(0, 1)
    
    # 添加数值标签
    for i, score in enumerate(scores):
        plt.text(i, score + 0.01, f'{score:.3f}', ha='center')
    
    # 模型类型信息
    plt.subplot(1, 3, 3)
    plt.text(0.1, 0.8, f"模型类型: {training_results['model_type']}", fontsize=12)
    plt.text(0.1, 0.6, f"训练准确率: {training_results['train_accuracy']:.3f}", fontsize=12)
    plt.text(0.1, 0.4, f"验证准确率: {training_results['validation_accuracy']:.3f}", fontsize=12)
    plt.text(0.1, 0.2, f"交叉验证: {training_results['cv_mean_score']:.3f} ± {training_results['cv_std_score']:.3f}", fontsize=12)
    plt.title('训练摘要')
    plt.axis('off')
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.show()
    
    print(f"训练结果图表已保存到: {save_path}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='训练语音情感识别模型')
    parser.add_argument('--data_path', type=str, required=True, help='数据集根目录路径')
    parser.add_argument('--model_type', type=str, default='random_forest', 
                       choices=['svm', 'knn', 'random_forest'], help='模型类型')
    parser.add_argument('--output_dir', type=str, default='models', help='模型输出目录')
    parser.add_argument('--plot_results', action='store_true', help='是否绘制训练结果图表')
    
    args = parser.parse_args()
    
    # 检查数据集路径
    if not os.path.exists(args.data_path):
        print(f"错误：数据集路径不存在: {args.data_path}")
        sys.exit(1)
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    try:
        # 加载数据集
        features, labels, files = load_dataset(args.data_path)
        
        # 初始化分类器
        print(f"\n开始训练 {args.model_type} 模型...")
        classifier = EmotionClassifier(model_type=args.model_type)
        
        # 训练模型
        training_results = classifier.train(features, labels)
        
        # 打印训练结果
        print(f"\n训练完成！")
        print(f"模型类型: {training_results['model_type']}")
        print(f"训练准确率: {training_results['train_accuracy']:.3f}")
        print(f"验证准确率: {training_results['validation_accuracy']:.3f}")
        print(f"交叉验证平均分: {training_results['cv_mean_score']:.3f} ± {training_results['cv_std_score']:.3f}")
        
        # 保存模型
        model_path = os.path.join(args.output_dir, f"emotion_classifier_{args.model_type}.joblib")
        scaler_path = os.path.join(args.output_dir, f"feature_scaler_{args.model_type}.joblib")
        
        classifier.save_model(model_path, scaler_path)
        
        # 保存训练报告
        report_path = os.path.join(args.output_dir, f"training_report_{args.model_type}.txt")
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write("语音情感识别模型训练报告\n")
            f.write("=" * 50 + "\n\n")
            f.write(f"模型类型: {training_results['model_type']}\n")
            f.write(f"训练准确率: {training_results['train_accuracy']:.3f}\n")
            f.write(f"验证准确率: {training_results['validation_accuracy']:.3f}\n")
            f.write(f"交叉验证: {training_results['cv_mean_score']:.3f} ± {training_results['cv_std_score']:.3f}\n\n")
            f.write("分类报告:\n")
            f.write(training_results['classification_report'])
        
        print(f"训练报告已保存到: {report_path}")
        
        # 绘制结果图表
        if args.plot_results:
            plot_path = os.path.join(args.output_dir, f"training_results_{args.model_type}.png")
            plot_training_results(training_results, plot_path)
        
        print("\n训练完成！")
        
    except Exception as e:
        print(f"训练过程中出错: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main() 